import random
import string
from app.settings import settings
import os

import fitz  # PyMuPDF
from random import choice
from string import ascii_letters, digits
from io import BytesIO
from fastapi import UploadFile
from typing import Union
from pathlib import Path


def generate_random_filename(extension):
    random_str = ''.join(random.choices(string.ascii_letters + string.digits, k=10))
    return f"{random_str}.{extension}"

def allowed_file(filename):
    ret = '.' in filename and filename.rsplit('.', 1)[1].lower() in settings.ALLOWED_IMAGES
    print(ret)
    return ret

def allowed_files(filename):
    ret = '.' in filename and filename.rsplit('.', 1)[1].lower() in settings.ALLOWED_EXTENSIONS
    print(ret)
    return ret

async def upload_image(file,path):
    if file.filename == '':
        return False
    if allowed_file(file.filename):
        extension = file.filename.rsplit('.', 1)[1].lower()
        random_filename = generate_random_filename(extension)
        url = os.path.join(settings.UPLOAD_FOLDER+"/"+path, random_filename)
    else:
        return False
    with open(url, "wb") as f:
        f.write(file.file.read())
    url = url.replace('\\', '/')
    return settings.APP_IP+url


async def upload_file(file: UploadFile, path: str) -> Union[str, bool]:
    if file.filename == '':
        return False

    if allowed_files(file.filename):
        extension = file.filename.rsplit('.', 1)[1].lower()
        random_filename = generate_random_filename(extension)

        # 如果是 PDF 文件，进行转换
        if extension == 'pdf':
            pdf_file_path = os.path.join(settings.UPLOAD_FOLDER, path, random_filename)
            # 保存 PDF 文件到临时路径
            with open(pdf_file_path, "wb") as f:
                f.write(file.file.read())

            # 转换 PDF 为 PNG
            return await convert_pdf_to_png(pdf_file_path, path)

        else:
            # 如果是其他类型的文件，直接保存
            url = os.path.join(settings.UPLOAD_FOLDER, path, random_filename)
            with open(url, "wb") as f:
                f.write(file.file.read())
            url = url.replace('\\', '/')
            return settings.APP_IP + url
    else:
        return False


async def convert_pdf_to_png(pdf_file_path: str, path: str) -> str:
    # 使用 PyMuPDF 转换 PDF 为 PNG
    doc = fitz.open(pdf_file_path)

    # 获取第一页
    page = doc.load_page(0)

    # 提高图像的清晰度，通过设置缩放因子
    zoom = 4.0  # 缩放因子，越大图像越清晰
    matrix = fitz.Matrix(zoom, zoom)  # 设置缩放矩阵
    pix = page.get_pixmap(matrix=matrix)

    # 生成随机 PNG 文件名
    random_filename = generate_random_filename('png')
    png_file_path = os.path.join(settings.UPLOAD_FOLDER, path, random_filename)

    # 保存图像为 PNG
    pix.save(png_file_path)

    # 删除原始 PDF 文件
    # os.remove(pdf_file_path)

    # 返回 PNG 文件的 URL
    png_url = png_file_path.replace('\\', '/')
    return settings.APP_IP + png_url

